package com.hj.aiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.redisson.api.RBucket;
import org.redisson.api.RedissonClient;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 基于Redis和Redisson的聊天记忆实现
 */
@Component
public class RedisChatMemory implements ChatMemory {

    private final RedissonClient redissonClient;
    private static final Kryo kryo = new Kryo();
    
    // Redis键前缀
    private static final String KEY_PREFIX = "chat:memory:";
    
    // 默认过期时间（7天）
    private static final Duration DEFAULT_TTL = Duration.ofDays(7);

    static {
        kryo.setRegistrationRequired(false);
        // 设置实例化策略
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    @Autowired
    public RedisChatMemory(RedissonClient redissonClient) {
        this.redissonClient = redissonClient;
    }

    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = buildKey(conversationId);
        RBucket<byte[]> bucket = redissonClient.getBucket(key);
        
        // 获取现有消息
        List<Message> conversationMessages = get(conversationId, Integer.MAX_VALUE);
        conversationMessages.addAll(messages);
        
        // 序列化并保存
        byte[] serializedData = serialize(conversationMessages);
        bucket.set(serializedData, DEFAULT_TTL);
    }

    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key = buildKey(conversationId);
        RBucket<byte[]> bucket = redissonClient.getBucket(key);
        
        byte[] data = bucket.get();
        if (data == null) {
            return new ArrayList<>();
        }
        
        List<Message> allMessages = deserialize(data);
        if (allMessages == null) {
            return new ArrayList<>();
        }
        
        // 返回最后N条消息
        return allMessages.stream()
                .skip(Math.max(0, allMessages.size() - lastN))
                .toList();
    }

    @Override
    public void clear(String conversationId) {
        String key = buildKey(conversationId);
        redissonClient.getBucket(key).delete();
    }
    
    /**
     * 设置对话过期时间
     */
    public void setExpiration(String conversationId, Duration ttl) {
        String key = buildKey(conversationId);
        RBucket<byte[]> bucket = redissonClient.getBucket(key);
        if (bucket.isExists()) {
            bucket.expire(ttl);
        }
    }
    
    /**
     * 获取对话剩余过期时间
     */
    public Duration getRemainingTTL(String conversationId) {
        String key = buildKey(conversationId);
        RBucket<byte[]> bucket = redissonClient.getBucket(key);
        if (bucket.isExists()) {
            return Duration.ofMillis(bucket.remainTimeToLive());
        }
        return Duration.ZERO;
    }
    
    /**
     * 检查对话是否存在
     */
    public boolean exists(String conversationId) {
        String key = buildKey(conversationId);
        return redissonClient.getBucket(key).isExists();
    }
    
    /**
     * 获取所有对话ID
     */
    public List<String> getAllConversationIds() {
        Iterable<String> keys = redissonClient.getKeys().getKeysByPattern(KEY_PREFIX + "*");
        List<String> conversationIds = new ArrayList<>();
        
        for (String key : keys) {
            conversationIds.add(key.substring(KEY_PREFIX.length()));
        }
        
        return conversationIds;
    }
    
    /**
     * 批量清理过期对话
     */
    public int cleanExpiredConversations() {
        Iterable<String> keys = redissonClient.getKeys().getKeysByPattern(KEY_PREFIX + "*");
        int cleanedCount = 0;
        
        for (String key : keys) {
            RBucket<byte[]> bucket = redissonClient.getBucket(key);
            if (bucket.isExists() && bucket.remainTimeToLive() == -1) {
                // 没有设置过期时间的键，设置默认过期时间
                bucket.expire(DEFAULT_TTL);
                cleanedCount++;
            }
        }
        
        return cleanedCount;
    }

    private String buildKey(String conversationId) {
        return KEY_PREFIX + conversationId;
    }

    private byte[] serialize(List<Message> messages) {
        try (ByteArrayOutputStream baos = new ByteArrayOutputStream();
             Output output = new Output(baos)) {
            kryo.writeObject(output, messages);
            output.flush();
            return baos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("Failed to serialize messages", e);
        }
    }

    @SuppressWarnings("unchecked")
    private List<Message> deserialize(byte[] data) {
        try (ByteArrayInputStream bais = new ByteArrayInputStream(data);
             Input input = new Input(bais)) {
            return kryo.readObject(input, ArrayList.class);
        } catch (IOException e) {
            throw new RuntimeException("Failed to deserialize messages", e);
        }
    }
}
